from typing import Optional, List
from collections import OrderedDict
import pyopencl as cl


def convert_bytes_to_readable(byte_n: int, round_n: int = 1) -> str:
    """ 将字节转为人类可读性更好的单位， 如GiB, Mib, Kib, Byte

    如 1025 被转换为 "1.0 Kib"

    :param byte_n: 原始字节数
    :param round_n: 转换后保留多少位小数
    :return: 转换后的字符串
    """
    one_kib = 1024
    one_mib = one_kib * 1024
    one_gib = one_mib * 1024

    for one_unit, name in zip([one_gib, one_mib, one_kib, 1], ['GiB', 'Mib', 'Kib', 'Byte']):
        if byte_n > one_unit:
            value = round(byte_n / one_unit, round_n)
            # 如果小数点后全为0, 则将value转为整数, 不显示小数点
            value = int(value) if int(value) == value else value
            return f'{value} {name}'


def get_device_type_name(device_type: int) -> str:
    device_types = {
        cl.device_type.CPU: 'CPU',
        cl.device_type.GPU: 'GPU',
        cl.device_type.ACCELERATOR: 'ACCELERATOR',
        cl.device_type.CUSTOM: 'CUSTOM',
    }

    return 'UNKNOWN' if device_type not in device_types else device_types[device_type]


def get_cache_type_name(cache_type: int) -> str:
    cache_types = {
        cl.device_mem_cache_type.READ_ONLY_CACHE: 'R',
        cl.device_mem_cache_type.READ_WRITE_CACHE: 'RW',
    }

    return 'UNKNOWN' if cache_type not in cache_types else cache_types[cache_type]


def work_size_to_str(work_size: List[int]) -> str:
    """ 工作项大小转为指定格式的字符串

    如 [512, 512, 512] 将被转换为 "512 x 512 x 512"

    :param work_size: 工作项大小
    :return: 转换后的字符串
    """
    result = ''
    for size in work_size:
        # 两个数值之间的乘号
        result += ' x ' if len(result) > 0 else ''
        result += str(size)
    return result


class DeviceInfo:
    def __init__(self, device: cl.Device):
        self._device = device
        self._platform = device.platform

        # 设备名称与平台相关
        self.name = device.name
        self.platform = self._platform.name
        self.vendor = device.vendor
        self.type = get_device_type_name(device.type)

        # 设备计算单元相关
        self.clock = f'{device.max_clock_frequency} MHZ' if device.max_clock_frequency != 0 else 'Unknown'
        self.cu = device.max_compute_units
        self.pe = device.preferred_work_group_size_multiple

        # 工作组和工作项相关
        self.max_work_group_size = device.max_work_group_size
        self.max_work_item_size = work_size_to_str(device.max_work_item_sizes)

        # 设备内存相关, 缓存无关
        self.global_size = convert_bytes_to_readable(device.global_mem_size)

        # 设备缓存相关
        self.global_cache = convert_bytes_to_readable(device.global_mem_cache_size)
        self.global_cache_type = get_cache_type_name(device.global_mem_cache_type)
        self.global_line = convert_bytes_to_readable(device.global_mem_cacheline_size)

    def __str__(self) -> str:
        result = {}
        for attr in dir(self):
            if '_' not in attr[0]:
                result[attr] = getattr(self, attr)
        result_sorted = OrderedDict(sorted(result.items()))
        return str(result_sorted)


class Devices:
    devices: Optional[List] = None

    def __init__(self) -> None:
        if Devices.devices is not None:
            return

        Devices.devices = []
        devices = Devices.devices
        for platform in cl.get_platforms():
            for device in platform.get_devices():
                device_info = DeviceInfo(device)
                devices.append(device_info)

    @staticmethod
    def get_device_info(index: int) -> DeviceInfo:
        if Devices.devices is None:
            Devices()
        if index > len(Devices.devices):
            # todo 完善报错信息
            raise IndexError
        return Devices.devices[index]


if __name__ == '__main__':
    for info in Devices.devices:
        print(info)
